import os
from urllib.parse import unquote

import requests.utils
from flask import Flask, request, Response

app = Flask(__name__)

# 定义逐跳头部列表
HOP_BY_HOP_HEADERS = frozenset([
    'connection', 'keep-alive', 'proxy-authenticate', 'proxy-authorization',
    'te', 'trailers', 'transfer-encoding', 'upgrade'
])


@app.route('/proxy', methods=['GET', 'POST', 'PUT', 'DELETE', 'PATCH', 'HEAD', 'OPTIONS'])
def proxy():
    # 从查询参数中获取目标URL
    url = request.args.get('url')

    if not url:
        return "Missing 'url' parameter", 400

    # URL解码
    decoded_url = unquote(url)

    try:
        # 准备请求头
        headers = {key: value for (key, value) in request.headers if key != 'Host'}

        # 准备请求数据
        data = request.get_data()

        # 发送请求到目标URL，保持原始方法和头部
        with requests.sessions.Session() as session:
            session.headers = requests.utils.CaseInsensitiveDict({})
            resp = session.request(
                method=request.method,
                url=decoded_url,
                headers=headers,
                data=data,
                stream=True,
                allow_redirects=False,
            )

        # 创建一个生成器，用于流式传输响应内容
        def generate():
            for chunk in resp.raw.stream(1024, decode_content=False):
                yield chunk

        # 创建一个新的Response对象
        proxied_resp = Response(generate(), status=resp.status_code)

        # 复制原始响应的头部信息
        for key, value in resp.headers.items():
            if key.lower() not in HOP_BY_HOP_HEADERS:  # 过滤响应头
                proxied_resp.headers[key] = value

        return proxied_resp

    except requests.RequestException as e:
        return f"Error fetching URL: {str(e)}", 525


if __name__ == '__main__':
    from waitress import serve

    address = os.getenv('ADDRESS', '127.0.0.1:7445')
    print(f"Starting proxy server on http://{address}")
    host, port = address.split(':')
    serve(app, host=host, port=int(port), threads=16)
